import argparse
import json
import os

import numpy as np
import matplotlib.pyplot as plt

from mpl_toolkits.mplot3d import Axes3D


def json2order(path):
    with open(path, 'r', encoding='utf-8') as f:
        dict_order = json.load(f)
    weather_dict = dict_order["attributes"]["weather"]
    scene_dict = dict_order["attributes"]["scene"]

    weather_order = [key for key, value in sorted(weather_dict.items(), key=lambda item: item[1], reverse=True)]
    scene_order = [key for key, value in sorted(scene_dict.items(), key=lambda item: item[1], reverse=True)]
    return weather_order, scene_order


def json2array(path, order_list):
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    weather_order, scene_order = order_list
    timeofday_order = {'daytime': 3, 'dawn/dusk': 2, 'night': 1, "undefined": 0}
    shape = (len(timeofday_order), len(weather_order), len(scene_order))
    result = np.zeros(shape=shape, dtype=np.uint32)

    for w_idx, weather in enumerate(weather_order):
        if weather in data['weather']:
            for s_idx, scene in enumerate(scene_order):
                if scene in data['weather'][weather]['scene']:
                    for timeofday, value in data['weather'][weather]['scene'][scene]['timeofday'].items():
                        t_idx = timeofday_order.get(timeofday, -1)
                        if t_idx != -1:  # 确保时间有效
                            result[t_idx, w_idx, s_idx] = value
    return result

def array2image(array477, order_list):
    weather, scene = order_list
    x = np.arange(len(weather))
    y = np.arange(len(scene))
    x, y = np.meshgrid(x, y)
    z = np.zeros_like(x, dtype=np.uint32)

    dx = dy = 0.5

    layers = len(array477)
    labels = ['Daytime', 'Dawn/Dusk', 'Night', 'Undefined']
    colors = ['red', 'green', 'blue', 'orange']

    # 创建三维图形
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')

    # 绘制堆积柱状图
    for i in range(layers):
        ax.bar3d(
            x.ravel(),
            y.ravel(),
            z.ravel(),
            dx,
            dy,
            array477[i].ravel(),
            color=colors[i],
            alpha=0.8,  # 设置透明度以便观察堆叠效果
        )
        z += array477[i]  # 更新 z 轴起点为当前层的顶部

    # 设置坐标轴标签
    ax.set_xticks(np.arange(len(weather)))
    ax.set_xticklabels(weather, fontsize=8)  # 设置 x 轴标签并旋转
    ax.set_yticks(np.arange(len(scene)))
    ax.set_yticklabels(scene, fontsize=8)  # 设置 y 轴标签

    # 设置标题和轴名称
    ax.set_title("Weather vs Scene with Stacked Z-axis", fontsize=14)
    ax.set_xlabel("Weather", fontsize=12)
    ax.set_ylabel("Scene", fontsize=12)
    ax.set_zlabel("Stacked Values", fontsize=12)

    # 手动创建图注
    handles = [plt.Rectangle((0, 0), 1, 1, color=colors[3-i], alpha=0.8) for i in range(layers)]
    ax.legend(handles, labels, title="Time of Day", loc="upper left", bbox_to_anchor=(1.05, 1))

    # 调整视角
    ax.view_init(elev=30, azim=45)  # 设置仰角和方位角
    ax.dist = 10  # 调整观察距离

    os.makedirs(opt.output_dir, exist_ok=True)
    output_path = os.path.join(opt.output_dir, "3D.png")
    plt.savefig(output_path)
    plt.tight_layout()
    plt.show()


if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input-path', type=str, default='./Json/7x7x4statistics.json', help='BDD format label path')
    parser.add_argument('--output-dir', type=str, default='./VisualizedImage', help='Output 3D cylindrical stacked image')
    parser.add_argument('--order-path', type=str, default='./Json/7+7+4statistics.json',
                        help='Get the approximate distribution to confirm the order on the axes')
    opt = parser.parse_args()

    order_path = opt.order_path
    order = json2order(order_path)

    input_path = opt.input_path
    array = json2array(input_path, order)

    array2image(array, order)
